import paddle.nn as nn
import paddle.nn.functional as F
import parl


class EasyGameModel(parl.Model):
    """ The Model for Atari env
    Args:
        obs_space (Box): observation space.
        act_space (Discrete): action space.
    """

    def __init__(self, obs_space, act_space):
        super(EasyGameModel, self).__init__()

        self.conv1 = nn.Conv2D(4, 32, 1, stride=1)
        # self.conv2 = nn.Conv2D(32, 64, 4, stride=2)
        # self.conv3 = nn.Conv2D(64, 64, 3, stride=1)

        self.flatten = nn.Flatten()
        self.fc = nn.Linear(288, 4)

        self.fc_pi = nn.Linear(4, act_space.n)
        self.fc_v = nn.Linear(4, 1)

    def value(self, obs):
        """ Get value network prediction
        Args:
            obs (np.array): current observation
        """
        obs = obs
        out = F.relu(self.conv1(obs))
        # out = F.relu(self.conv2(out))
        # out = F.relu(self.conv3(out))

        out = self.flatten(out)
        out = F.relu(self.fc(out))
        value = self.fc_v(out)
        return value

    def policy(self, obs):
        """ Get policy network prediction
        Args:
            obs (np.array): current observation
        """
        obs = obs
        out = F.relu(self.conv1(obs))
        # out = F.relu(self.conv2(out))
        # out = F.relu(self.conv3(out))

        out = self.flatten(out)
        out = F.relu(self.fc(out))
        logits = self.fc_pi(out)

        return logits
